package com.ics.atable.chat.service.impl;

import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.ics.atable.chat.common.milvus.MilvusMultiVectorStore;
import com.ics.atable.chat.constant.VectorDatabaseConstant;
import com.ics.atable.chat.expection.BusinessException;
import com.ics.atable.chat.model.dto.ItemRetrievalDTO;
import com.ics.atable.chat.model.dto.MatchedItem;
import com.ics.atable.chat.service.EmbeddingService;
import com.ics.atable.chat.service.RetrievalTestService;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static com.ics.atable.chat.constant.VectorDatabaseConstant.ITEM_COLLECTION_NAME;
import static com.ics.atable.chat.expection.ErrorCode.OPERATION_ERROR;
import static com.ics.atable.chat.expection.ErrorCode.PARAMS_ERROR;

/**
 * @author haohongbin
 */
@Slf4j
@Service
public class RetrievalTestServiceImpl implements RetrievalTestService {

    private static final String itemCollectionName = ITEM_COLLECTION_NAME;

    private final EmbeddingService embeddingService;

    private final MilvusMultiVectorStore milvusMultiVectorStore;

    public RetrievalTestServiceImpl(EmbeddingService embeddingService, MilvusMultiVectorStore milvusMultiVectorStore) {
        this.embeddingService = embeddingService;
        this.milvusMultiVectorStore = milvusMultiVectorStore;
    }

    @Override
    public List<MatchedItem> itemRetrieval(ItemRetrievalDTO itemRetrievalDTO) {
        try {
            List<MatchedItem> matchedItems = new ArrayList<>();

            EmbeddingResponse embeddingResponse = embeddingService.embedding(List.of(itemRetrievalDTO.getContent()));
            float[] vector = embeddingResponse.getResults().getFirst().getOutput();

            // 检索 item
            milvusMultiVectorStore.loadCollection(itemCollectionName);
            List<String> denseFieldNames = new ArrayList<>();
            denseFieldNames.add(VectorDatabaseConstant.ItemField.ITEM_NAME.getCode());
            denseFieldNames.add(VectorDatabaseConstant.ItemField.ITEM_FATHER_NAME.getCode());
            denseFieldNames.add(VectorDatabaseConstant.ItemField.ITEM_GRANDFATHER_NAME.getCode());

            if (denseFieldNames.size() != itemRetrievalDTO.getWeights().size()) {
                throw new BusinessException(PARAMS_ERROR, "权重数量与字段数量不一致");
            }

            String itemFilter = ""; // 过滤条件

            List<BaseVector> queryVectors = List.of(new FloatVec(vector));

            SearchResp searchResp = milvusMultiVectorStore.hybridSearch(
                    itemCollectionName,
                    queryVectors,
                    denseFieldNames,
                    itemRetrievalDTO.getWeights(),
                    itemRetrievalDTO.getTopK(),
                    itemFilter
            );
            matchedItems = searchRespToItem(searchResp);

            return matchedItems;

        } catch (Exception e) {
            log.error("检索项失败: {}", e.getMessage(), e);
            throw new BusinessException(OPERATION_ERROR, "检索项失败");
        }  finally {
            milvusMultiVectorStore.releaseCollection(itemCollectionName);
        }
    }

    /**
     * 将搜索结果转换为Item对象列表
     *
     * @param searchResp 搜索结果
     * @return Item对象列表
     */
    private List<MatchedItem> searchRespToItem(SearchResp searchResp) {
        List<MatchedItem> matchedItems = new ArrayList<>();
        for (List<SearchResp.SearchResult> searchResultList : searchResp.getSearchResults()) {
            for (SearchResp.SearchResult searchResult : searchResultList) {
                MatchedItem matchedItem = new MatchedItem();
                matchedItem.setItemUuid(searchResult.getEntity().get("item_uuid").toString());

                // 从metadata中提取item_name和item_father_name
                Object metadataObj = searchResult.getEntity().get("metadata");
                if (metadataObj != null) {
                    String metadataStr = metadataObj.toString();

                    try {
                        // 修复转义问题：尝试处理可能被转义的JSON字符串
                        String processedMetadataStr = metadataStr;

                        // 检查是否是被引号包围的JSON字符串
                        if (metadataStr.startsWith("\"{") && metadataStr.endsWith("}\"")) {
                            // 去除首尾引号并解转义
                            processedMetadataStr = metadataStr.substring(1, metadataStr.length() - 1)
                                    .replace("\\\"", "\"")
                                    .replace("\\n", "\n")
                                    .replace("\\\\", "\\");
                        }

                        JsonObject metadataJson = JsonParser.parseString(processedMetadataStr).getAsJsonObject();
                        matchedItem.setItemName(metadataJson.has("item_name") ?
                                metadataJson.get("item_full_name").getAsString() : "");
                        String fullName = metadataJson.has("item_full_name") ?
                                metadataJson.get("item_full_name").getAsString() : "";
                        log.info("fullName: {}", fullName);
                        List<String> fullNameList = Arrays.asList(fullName.split("-"));
                        if (fullNameList.size() == 1) {
                            matchedItem.setItemName(fullName);
                        } else if (fullNameList.size() == 2) {
                            matchedItem.setItemName(fullName);
                            matchedItem.setItemFatherName(fullNameList.get(1));
                        } else if (fullNameList.size() >= 3) {
                            matchedItem.setItemName(fullName);
                            matchedItem.setItemFatherName(fullNameList.get(1));
                            matchedItem.setItemGrandfatherName(fullNameList.get(2));
                        }
                    } catch (Exception e) {
                        log.warn("解析metadata JSON失败: {}", metadataStr, e);
                        matchedItem.setItemName("");
                        matchedItem.setItemFatherName("");

                    }
                } else {
                    matchedItem.setItemName("");
                    matchedItem.setItemFatherName("");
                }

                matchedItem.setScore(searchResult.getScore());
                matchedItems.add(matchedItem);
            }
        }
        return matchedItems;
    }

    /**
     * 解析item_list JSON数组为字符串列表
     * @param jsonArray JSON数组
     * @return 字符串列表
     */
    private List<String> parseItemList(JsonArray jsonArray) {
        List<String> itemList = new ArrayList<>();
        try {
            for (JsonElement element : jsonArray) {
                if (element.isJsonPrimitive()) {
                    itemList.add(element.getAsString());
                }
            }
        } catch (Exception e) {
            log.warn("解析item_list数组失败: {}", jsonArray, e);
        }
        return itemList;
    }
}
